import json
import sys
from pathlib import Path
from collections import Counter

EXPECTED_THEMES = [
    "乡村生活类",
    "农村美食类",
    "农技展示类",
    "民间文化类",
    "农事劳作类",
    "电商助农类",
    "剧情情感类",
]


def read_theme(path: Path) -> str:
    try:
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.startswith('视频主题:'):
                    return line.split(':', 1)[1].strip()
    except Exception:
        return ""
    return ""


def summarize_themes(results_dir: Path):
    files = sorted(results_dir.glob('*.txt'))
    if not files:
        print(f"[错误] 未找到TXT: {results_dir}")
        sys.exit(1)

    dist = Counter()
    missing = []
    extras = Counter()

    for p in files:
        theme = read_theme(p)
        if not theme:
            missing.append(p.name)
            continue
        if theme in EXPECTED_THEMES:
            dist[theme] += 1
        else:
            extras[theme] += 1

    total = len(files)
    parsed = sum(dist.values()) + sum(extras.values())

    present_expected = [t for t in EXPECTED_THEMES if dist.get(t, 0) > 0]
    missing_expected = [t for t in EXPECTED_THEMES if dist.get(t, 0) == 0]
    extra_categories = [k for k in extras.keys()]

    # 输出控制台概要
    print(f"[总计] TXT 文件数: {total}")
    print(f"[解析成功] {parsed}，[缺失主题] {len(missing)}")
    print("[视频主题分布]")
    for t in EXPECTED_THEMES:
        print(f"  {t} -> {dist.get(t, 0)}")
    if extras:
        print("[其他类别]")
        for k, v in extras.most_common():
            print(f"  {k} -> {v}")
    print("[是否覆盖七类] {}".format("是" if len(missing_expected) == 0 else "否"))

    data = {
        "total_files": total,
        "parsed": parsed,
        "missing_theme_files": missing,
        "distribution": dist,
        "expected_categories": EXPECTED_THEMES,
        "present_expected": present_expected,
        "missing_expected": missing_expected,
        "extra_categories": {k: v for k, v in extras.items()},
    }
    # 将 Counter 转普通 dict
    data["distribution"] = {k: v for k, v in dist.items()}

    out_path = results_dir.parent / "theme_counts.json"
    try:
        with open(out_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print(f"[JSON已写入] {out_path}")
    except Exception as e:
        print(f"[写入JSON失败] {e}")


if __name__ == '__main__':
    results_dir = Path('results').resolve()
    summarize_themes(results_dir)
